from fit_py_ask_number_demo.configs.knowledge_base import ENUM_BASE, SYNONYM_BASE, RULE_BASE
from fit_py_ask_number_demo.client import llm_client
from fit_py_ask_number_demo.configs.prompts import info_extract_prompt


class InfoExtractor:
    def __init__(self, query):
        self.query = query
        self.extracted_words = []
        self.matched_field_info = {}
        self.ambiguous_words = {}

    def run(self):
        self.extracted_words = []
        self.matched_field_info = {}
        self.process_enum()
        self.process_rule()
        return self.matched_field_info, self.ambiguous_words

    def run_synonym_replace(self):
        self.process_synonym()
        return self.query

    def process_synonym(self):
        for standard_word, synonyms_list in SYNONYM_BASE.items():
            for synonym in synonyms_list:
                # 检查标准词是否已经包含在文本中
                if synonym in self.query and standard_word not in self.query:
                    self.query = self.query.replace(synonym, standard_word)

    def process_enum(self):
        # 遍历每个数据库字段和对应的枚举值
        for field, values in ENUM_BASE.items():
            for value in values:
                if value in self.query:
                    # 检查是否有更短的值已经被匹配，有的话需要删除
                    delete_words = []
                    for word in self.extracted_words:
                        if word in value and len(word) < len(value):
                            delete_words.append(word)
                        elif word == value:
                            if self.ambiguous_words.get(word):
                                self.ambiguous_words[word].append({field: {"in": [value]}})
                            else:
                                self.ambiguous_words[word] = [{field: {"in": [value]}}, self.matched_field_info[value]]

                    if delete_words:
                        new_extracted_words = []
                        for word in self.extracted_words:
                            if word not in delete_words:
                                new_extracted_words.append(word)
                        self.extracted_words = new_extracted_words

                        for word in delete_words:
                            self.matched_field_info.pop(word)

                    self.matched_field_info[value] = {field: {"in": [value]}}
                    self.extracted_words.append(value)

    def process_rule(self):
        # 遍历每个数据库字段和对应的枚举值
        for value, rule in RULE_BASE.items():
            if value in self.query:
                # 检查是否有更短的值已经被匹配，有的话需要删除
                delete_words = []
                for word in self.extracted_words:
                    if word in value and len(word) < len(value):
                        delete_words.append(word)
                    elif word == value:
                        if self.ambiguous_words.get(word):
                            self.ambiguous_words[word].append(rule)
                        else:
                            self.ambiguous_words[word] = [rule, self.matched_field_info[value]]

                if delete_words:
                    new_extracted_words = []
                    for word in self.extracted_words:
                        if word not in delete_words:
                            new_extracted_words.append(word)
                    self.extracted_words = new_extracted_words

                    for word in delete_words:
                        self.matched_field_info.pop(word)
                self.matched_field_info[value] = rule
                self.extracted_words.append(value)

    def extract_info_by_llm(self):
        user_content = f"<Question>\n{self.query}"
        semantic_info = llm_client.deepseek_client(info_extract_prompt, user_content)
        return semantic_info


def synonym_replace_impl(query):
    info_extractor = InfoExtractor(query)
    new_query = info_extractor.run_synonym_replace()
    return new_query


def field_info_extract_impl(query):
    info_extractor = InfoExtractor(query)
    field_info, ambiguous_info = info_extractor.run()
    return str(field_info), ambiguous_info


def info_extract_impl(query):
    new_query = synonym_replace_impl(query)
    field_info, ambiguous_info = field_info_extract_impl(new_query)
    semantic_info = InfoExtractor(query).extract_info_by_llm()
    return new_query, field_info, semantic_info, ambiguous_info


if __name__ == '__main__':
    query = "24年战略地区东进计划北京市KPI收入"
    res = info_extract_impl(query)
    print(res)
